# -*- coding: utf-8 -*-
# @Time    : 2019/2/22 10:47
# @Author  : cj
# @File    : annotation.py
import os
import xml.etree.cElementTree as ET
from xml.etree.ElementTree import Element, SubElement, ElementTree

from dataset.annotation_object import AnnotationObject
from dataset.bounding_box import BoundingBox
from dataset.point import Point


class Annotation:
    """
    类Annotation保存一个完成的xml信息
    """

    def __init__(self, image_filename=None, image_size=None, annotation_objects=None):
        self._image_filename = image_filename
        # 图片大小保存为元组[height,width,3]
        self._image_size = image_size
        self._annotation_objects = annotation_objects

    def get_boundingboxes(self):
        """
        :return:返回一个标注项的所有boundingbox:list
        """
        boundingboxes = []
        for annotation_object in self._annotation_objects:
            boundingboxes.append(annotation_object.get_boundingbox())
        return boundingboxes

    def set_boundingboxes(self, boundingboxes):
        """
        设置boundingboxes
        :param boundingboxes: 设置的boundingboxes要与_annotation_objects的每一个对应起来
        :return:
        """
        if len(boundingboxes) != len(self._annotation_objects):
            raise ValueError("置的boundingboxes与_annotation_objects不对应")
        for idx, annotation_object in enumerate(self._annotation_objects):
            boundingboxes.append(annotation_object.set_boundingbox(boundingboxes[idx]))

    def get_classnames(self):
        """
        :return: 返回一个标注项的所有classname：list,与boundingboxes对应
        """
        classnames = []
        for annotation_object in self._annotation_objects:
            classnames.append(annotation_object.get_class_name())
        return classnames

    def load_from_xml_file(self, xml_file_path: str):
        """
         加载xml文件，将xml里的信息保存好
        """

        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        image_filename = root.find('filename').text
        self._image_filename = image_filename

        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)
        self._image_size = [height, width, 3]

        annotation_objects = []
        for obj in root.iter('object'):
            class_name = obj.find('./name').text
            difficult = bool(int(obj.find('difficult').text))
            xml_box = obj.find('bndbox')
            _xmin = int(xml_box.find('xmin').text)
            _xmax = int(xml_box.find('xmax').text)
            _ymin = int(xml_box.find('ymin').text)
            _ymax = int(xml_box.find('ymax').text)
            annotation_objects.append(
                AnnotationObject(class_name, difficult, BoundingBox(Point(_xmin, _ymin), Point(_xmax, _ymax))))
        self._annotation_objects = annotation_objects

    def save_to_xml_file(self, path: str, database1: str):
        """
        函数功能:保存xml到磁盘
        :param path:要保存xml的路径
        :param database1: 数据库来源
        """
        # 生成根节点
        annotation = Element('annotation')
        # 生成第一个子节点 filename
        filename = SubElement(annotation, 'filename')
        filename.text = self._image_filename
        # 生成第二个子节点 source
        source = SubElement(annotation, 'source')
        # source 节点的子节点
        database = SubElement(source, 'database')
        database.text = database1
        # 生成 root 的第三个子节点 size
        size = SubElement(annotation, 'size')
        # size 节点的子节点
        width = SubElement(size, 'width')
        height = SubElement(size, 'height')
        depth = SubElement(size, 'depth')
        height.text = str(self._image_size[0])
        width.text = str(self._image_size[1])
        depth.text = str(self._image_size[2])

        # 生成 root 的object子节点
        for annotation_object in self._annotation_objects:
            object = SubElement(annotation, 'object')
            # object 节点的子节点
            name = SubElement(object, 'name')
            name.text = annotation_object.get_class_name()
            difficult = SubElement(object, 'difficult')
            difficult.text = str(annotation_object.is_difficult())

            bndbox = SubElement(object, 'bndbox')
            xmin = SubElement(bndbox, 'xmin')
            xmax = SubElement(bndbox, 'xmax')
            ymin = SubElement(bndbox, 'ymin')
            ymax = SubElement(bndbox, 'ymax')

            boundingbox = annotation_object.get_boundingbox()

            upper_left = boundingbox.get_upper_left()
            bottom_right = boundingbox.get_bottom_right()
            xmin.text = str(upper_left.x)
            ymin.text = str(upper_left.y)
            xmax.text = str(bottom_right.x)
            ymax.text = str(bottom_right.y)

        tree = ElementTree(annotation)
        image_filename = self._image_filename
        xml_filename = image_filename.split('.')[0:-1][0] + '.xml'
        tree.write(os.path.join(path, xml_filename), encoding='utf-8')

    def get_image_filename(self):
        """
        :return: 返回图片文件名
        """
        return self._image_filename

    def get_image_size(self):
        """
        :return:返回图片大小
        """
        return self._image_size

    def get_annotation_objects(self):
        """
        :return: 返回所有的AnnotationObject对象
        """
        return self._annotation_objects
